"""
This part is based on the dataset class implemented by pytorch,
including train_dataset and test_dataset, as well as data augmentation
"""
from numpy.lib.function_base import angle
from torch.utils.data import Dataset
import torch
import numpy as np
import random
import torch.nn.functional as F
from torchvision import transforms
from torchvision.transforms.functional import normalize,center_crop,rotate,affine,adjust_brightness, adjust_contrast, adjust_gamma,adjust_saturation,hflip,vflip

# https://pytorch.org/vision/master/_modules/torchvision/transforms/functional.html
# rotate affine   adjust_brightness adjust_contrast adjust_gamma 

class RandomFlip:
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, mask=None):
        if random.uniform(0,1) < self.p:
            img = hflip(img)
            if mask is not None:   
                mask = hflip(mask)


        if random.uniform(0,1) < self.p:
            img = vflip(img)
            if mask is not None:   
                mask = vflip(mask)
        if mask is not None:
            return img,mask
        return img
        
class RandomRotate:
    def __init__(self, degrees):
        self.degrees = degrees



    def __call__(self, img, mask=None):
        angle = random.uniform(-self.degrees, self.degrees)
        if mask is not None:
            return rotate(img, angle), rotate(mask, angle)
        else:
            return rotate(img, angle)


class Center_Crop:
    def __init__(self, size):
        self.size = size

    def __call__(self, img, label=None):
        if label is not None:
            crop_img = center_crop(img,self.size)
            crop_label = center_crop(label,self.size)
            return crop_img, crop_label
        else:
            crop_img = center_crop(img,self.size)
            return crop_img

class ToTensor:
    def __init__(self):
        self.to_tensor = transforms.ToTensor()

    def __call__(self, img, mask=None):
        if mask is not None:
            img = self.to_tensor(img)
            mask =self.to_tensor(mask)
            return img, mask
        else:
            img = self.to_tensor(img)
            return img

class Normalize:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, img, mask=None):
        if mask is not None:
            return normalize(img, self.mean, self.std, False), mask
        else:
            return normalize(img, self.mean, self.std, False)


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img, mask=None):
        if mask is not None:
            for t in self.transforms:
                img, mask = t(img, mask)
            return img, mask
        else:
            for t in self.transforms:
                img = t(img)
            return img

class RandomAffine:
    def __init__(self,degrees=10, translate=(0.1,0.1),scale=(0.9,1.1)):
        self.degrees=degrees
        self.translate=translate
        self.scale = scale

    def __call__(self, img, mask=None):
        translate_x = int(img.shape[0] * self.translate[0])
        translate_y = int(img.shape[1] * self.translate[1])

        angle = random.uniform(-self.degrees, self.degrees)
        translate_tuple = (random.randint(-translate_x,translate_x),random.randint(-translate_y,translate_y))


        scale_ = random.uniform(self.scale[0],self.scale[1])
        
        affine_img = affine(img,angle,translate_tuple,scale=scale_ ,shear=(0,0))
        if mask is not None:
            affine_mask = affine(mask,angle,translate_tuple,scale=scale_ ,shear=(0,0))

            return affine_img,affine_mask
        else:
            return affine_img
class ColorJitter:
    def __init__ (self,brightness=0.1,contrast=0.1,saturation=0.1):
        self.brightness = brightness
        self.contrast = contrast
        self.saturation = saturation
    def __call__(self, img, mask=None):
        brightness = random.uniform(1-self.brightness,1+self.brightness)
        img = adjust_brightness(img,brightness)
        # mask = adjust_brightness(mask,brightness)
        
        contrast = random.uniform(1-self.contrast,1+self.contrast)
        img = adjust_contrast(img,contrast)
        # mask = adjust_contrast(mask,contrast)

        saturation = random.uniform(1-self.saturation,1+self.saturation)
        img = adjust_saturation(img,saturation)
        # mask = adjust_saturation(mask,saturation)
        if mask is not None:
            return img,mask
        else:
            return img

